# indicators/indicator_calculator.py
import pandas as pd

from data.mysql_db import DataManager
from utils.MyTT import KDJ, BBI, MACD


class TechnicalIndicators:
    def __init__(self, data: pd.DataFrame):
        self._data = data.copy()

    def calculate_bbi(self, periods=None) -> pd.DataFrame:
        """
        计算多空指数 (BBI)

        参数:
            periods (dict{'M1': 5, 'M2': 10, 'M3': 20, 'M4': 30}): 均线周期列表

        返回:
            DataFrame: 包含BBI指标的DataFrame
        """
        if periods is None:
            periods = {'M1': 5, 'M2': 10, 'M3': 20, 'M4': 30}

        # 一次性计算所有均线并取平均
        self._data['BBI'] = BBI(self._data['close'], **periods)
        # 四舍五入
        self._data['BBI'] = self._data['BBI'].round(2)
        # 选择需要的列（直接返回 DataFrame 切片）
        return self._data[['trade_date', 'close', 'BBI']]

    def calculate_kdj(self, n=9, m1=3, m2=3) -> pd.DataFrame:
        """
        计算随机指标 (KDJ)

        参数:
            n (int): RSV周期
            m1 (int): K平滑系数
            m2 (int): D平滑系数

        返回:
            DataFrame: 包含KDJ指标的DataFrame
        """

        self._data['K'], self._data['D'], self._data['J'] = KDJ(CLOSE=self._data["close"], HIGH=self._data["high"],
                                                                LOW=self._data["low"], N=n, M1=m1, M2=m2)
        self._data['K'] = self._data['K'].round(2)
        self._data['D'] = self._data['D'].round(2)
        self._data['J'] = self._data['J'].round(2)

        return self._data[['trade_date', 'close', 'K', 'D', 'J']]

    def calculate_macd(self, short=12, long=26, signal=9) -> pd.DataFrame:
        """
        计算异同移动平均线 (MACD)

        参数:
            short (int): 短期EMA周期
            long (int): 长期EMA周期
            signal (int): 信号线周期

        返回:
            DataFrame: 包含MACD指标的DataFrame
        """

        self._data['DIF'], self._data['DEA'], self._data['MACD'] = MACD(self._data['close'], short, long, signal)
        return self._data[['trade_date', 'close', 'DIF', 'DEA', 'MACD']]

    @property
    def data(self):
        return self._data


# 示例用法
if __name__ == '__main__':
    data_manager = DataManager()
    ts_code = '600610.SH'
    df = data_manager.get_stock_daily_data(ts_code)
    calculator = TechnicalIndicators(df)
    calculator.calculate_bbi()
    calculator.calculate_kdj()
    calculator.calculate_macd()
    print(calculator.data.tail(5))
